import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import classification_report, accuracy_score
from sklearn.svm import SVC
from joblib import dump

class classIfcation:
    def __init__(self):
        self.df = pd.read_csv('fina_indicator.csv')
        
    def get_condition(self):
        df = self.df
        df['max_ratio'] = df['max_close'] / df['the_close']
        df['min_ratio'] = df['min_close'] / df['the_close']
        
        high_return_threshold =df['max_ratio'] .quantile(0.4)
        high_risk_threshold =df['min_ratio'] .quantile(0.4)
        
        condition = [
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] < high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] < high_risk_threshold)
        ]
        choices = ['高收益高风险','低收益高风险','高收益低风险','低收益低风险']
        df['category'] = np.select(condition, choices, default='未知')
        
        features = df[['eps','total_revenue_ps','undist_profit_ps', 'gross_margin', 'fcff','fcfe','tangible_asset','bps','grossprofit_margin','npta']]
        
        label_encoder = LabelEncoder()
        df['category_encoded'] = label_encoder.fit_transform(df['category'])
        
        scaler = StandardScaler()
        X = scaler.fit_transform(features)
        y = df['category_encoded']
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        dump(scaler, 'scaler.joblib')
        return X_train, X_test, y_train, y_test, label_encoder

    def knn_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        knn = KNeighborsClassifier(n_neighbors=3)
        knn.fit(X_train, y_train)
        y_pred = knn.predict(X_test)
        print(classification_report(y_test, y_pred, target_names=label_encoder.classes_))
        dump(knn, 'knn_model.joblib')
        dump(label_encoder, 'label_encoder.joblib')

    def svc_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        svc = SVC(kernel='linear', C=1.0, random_state=42)
        svc.fit(X_train, y_train)
        y_pred = svc.predict(X_test)
        print(classification_report(y_test, y_pred, target_names=label_encoder.classes_))
        dump(svc, 'svc_model.joblib')
        dump(label_encoder, 'label_encoder.joblib')
        


if __name__ == '__main__':
    ci = classIfcation()
    X_train, X_test, y_train, y_test, label_encoder = ci.get_condition()
    #ci.knn_utils(X_train, X_test, y_train, y_test, label_encoder)
    ci.svc_utils(X_train, X_test, y_train, y_test, label_encoder)